package org.ala.tiktools.sharddb.mybatis.executor;

import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.stream.Collectors;

import org.ala.tiktools.logs.LoggerFactory;
import org.ala.tiktools.mybatis.spring.TikToolsMyBatisSpringContext;
import org.ala.tiktools.sharddb.db.ThreadLocalLookupKey;
import org.ala.tiktools.sharddb.exception.ShardDBRuntimeException;
import org.ala.tiktools.sharddb.mybatis.executor.ShardExecutor.Result;
import org.ala.tiktools.sharddb.mybatis.interceptor.Shard;
import org.ala.tiktools.sharddb.mybatis.interceptor.ShardInfoThreadLocal;
import org.ala.tiktools.sharddb.mybatis.router.ShardInfo;
import org.ala.tiktools.sharddb.mybatis.transaction.ConcurrentMannerEnums;
import org.ala.tiktools.sharddb.mybatis.transaction.IShardCompensation;
import org.ala.tiktools.sharddb.mybatis.transaction.ShardTransaction;
import org.ala.tiktools.sharddb.tx.IShardTransaction;
import org.ala.tiktools.sharddb.tx.ShardDataSourceTransactionManager;
import org.ala.tiktools.sharddb.tx.ShardTransactionThreadLocal;
import org.ala.tiktools.tools.CachingKVTools;
import org.slf4j.Logger;
import org.springframework.util.Assert;

/**
 * 事务执行器
 * <p>  逻辑比较多，单独抽出来
 *
 * @author ala
 * @date 2025-02-24 13:50
 */
public class ShardTxExecutor {


    static Logger log = LoggerFactory.getLogger("shard", "tx");


    /**
     *  执行事务方法
     */
    public Map<ShardInfo, ShardExecutor.Result> doShardInvokeWithTx(Shard anno, ShardTransaction txAnno, ThreadPoolExecutor threadPool, Map<ShardInfo, Object[]> shardInfoMap, Object proxy, Method method) {
        //  如果是强一致提交，直接走分布式事务
        if (txAnno.strongConsistency()) {
            return strongConsistency(anno, txAnno, threadPool, shardInfoMap, proxy, method);
        }
        //  否则挑出成功和失败的，自行决定是回滚还是最终一致掉
        else {
            return eventualConsistency(anno, txAnno, threadPool, shardInfoMap, proxy, method);
        }
    }


    /**
     *  最终一致方案
     */
    protected Map<ShardInfo, ShardExecutor.Result> eventualConsistency(Shard anno, ShardTransaction txAnno, ThreadPoolExecutor threadPool, Map<ShardInfo, Object[]> shardInfoMap, Object proxy, Method method) {
        Map<ShardInfo, ShardExecutor.Result> res = new ConcurrentHashMap<>();
        //  事务补偿
        IShardCompensation compensation = compensationMapper.get(method);
        Assert.notNull(compensation, String.format("最终一致时，必须提供补偿机制 method:%s class:%s", method.getName(), method.getDeclaringClass().getSimpleName()));

        //  直接跑事务，记录成功失败的
        CountDownLatch latch = new CountDownLatch(shardInfoMap.size());
        long now = System.currentTimeMillis();
        shardInfoMap.forEach((shardInfo, args) -> {
            try {
                threadPool.execute(() -> {
                    try {
                        ThreadLocalLookupKey.setLookup(shardInfo.getShardDB());
                        res.put(shardInfo, ShardExecutor.doInvoke(anno, shardInfo, method, args).timeOfExecute(System.currentTimeMillis() - now));
                    } finally {
                        ThreadLocalLookupKey.clear();
                        latch.countDown();
                    }
                });
            } catch (Exception e){
                res.put(shardInfo, ShardExecutor.Result.error(shardInfo, e).timeOfExecute(System.currentTimeMillis() - now));
                latch.countDown();
            }
        });

        //  等待所有人执行完成，没执行完的判为超时
        try {
            latch.await(txAnno.timeout(), TimeUnit.MILLISECONDS);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        //  执行成功的和执行失败的
        List<ShardExecutor.Result> success = new ArrayList<>(), errors = new ArrayList<>();
        shardInfoMap.forEach((shardInfo, args) -> {
            ShardExecutor.Result r = res.get(shardInfo);
            if (r == null) {
                r = ShardExecutor.Result.error(shardInfo, new TimeoutException()).timeOfExecute(System.currentTimeMillis() - now);
                res.put(shardInfo, r);
                errors.add(r);
            } else if (!r.isSuccess()) {
                errors.add(r);
            } else {
                success.add(r);
            }
        });

        //  挑出成功失败的，执行补偿
        if (!errors.isEmpty()) {
            try {
                compensation.doCompensation(success, errors);
            } catch (Exception e) {
                log.error(String.format("执行事务补偿失败 method:%s class:%s", method.getName(), method.getDeclaringClass().getSimpleName()), e);
            }
        }

        return res;
    }


    /**
     *  强一致方案
     */
    protected Map<ShardInfo, ShardExecutor.Result> strongConsistency(Shard anno, ShardTransaction txAnno, ThreadPoolExecutor threadPool, Map<ShardInfo, Object[]> shardInfoMap, Object proxy, Method method) {
        //  拿事务（如果事务对象是从线程环境中拿出来的）
        IShardTransaction tx = ShardTransactionThreadLocal.getTx();
        boolean commit = false;
        if (tx == null) {
            tx = getShardDataSourceTransactionManager().open(txAnno.timeout(), threadPool);
            commit = true;
        }

        Map<ShardInfo, ShardExecutor.Result> res = new ConcurrentHashMap<>();
        //  事务补偿
        IShardCompensation compensation = compensationMapper.get(method);

        //  事务开始计时
        long now = System.currentTimeMillis();
    	if (ConcurrentMannerEnums.Parallel_Shard_Table.equals(txAnno.concurrent())) {
            res = parallelShardTable(anno, tx, shardInfoMap, proxy, method);
        } else {
            res = parallelShardDB(anno, tx, shardInfoMap, proxy, method);
        }
        

        //  提交事务（任何子事务执行过程中抛出异常均回滚）
        if (commit) {
            try {
                tx.commit((sl, fl) -> {
                    if (compensation != null) {
                        compensation.doCompensation(
                                sl.stream().map(t -> ((ShardExecutor.Result) t.getMetadata()).success(true).res(t.getResult()).timeOfExecute(System.currentTimeMillis() - now)).collect(Collectors.toList()),
                                fl.stream().map(t -> ((ShardExecutor.Result) t.getMetadata()).success(false).e(t.getE()).timeOfExecute(System.currentTimeMillis() - now)).collect(Collectors.toList())
                        );
                    }
                });
            } catch (Exception e) {
                throw new ShardDBRuntimeException(String.format("本地分布式事务执行异常，已回滚 method:%s class:%s", method.getName(), method.getDeclaringClass().getSimpleName()), e);
            }
        }

        return res;
    }
    /**
     *  分库并行
     */
    protected Map<ShardInfo, ShardExecutor.Result> parallelShardDB(Shard anno, IShardTransaction tx, Map<ShardInfo, Object[]> shardInfoMap, Object proxy, Method method) {
        Map<ShardInfo, ShardExecutor.Result> res = new HashMap<>();

        //  shardInfoMap按照分库再聚合起来
        Map<String, Map<ShardInfo, Object[]>> shardInfoByDB = new HashMap<>();
        shardInfoMap.forEach((s,o) -> {
            shardInfoByDB.computeIfAbsent(s.getShardDB(), k -> new HashMap<>()).put(s, o);
        });

        //  每个分库下串行执行
        shardInfoByDB.forEach((db, m) -> {
            String name = String.format("shard_tx_%s_%s_%s", method.getDeclaringClass().getSimpleName(), method.getName(), db);
            try {
                tx.withTx2PCAsync(name, db, null, () -> {
                    //  分库下的所有分表在一个线程中串行执行
                    m.forEach((shardInfo, shardArgs) -> {
//                        ShardInfoThreadLocal.set(shardInfo);
                        try {
                        	Result r = ShardExecutor.doInvoke(anno, shardInfo, method, shardArgs);
                        	if (!r.isSuccess()) {
                        		throw new ShardDBRuntimeException(r.getE());
                        	}
                        	res.put(shardInfo, r);
                        } finally {
//                            ShardInfoThreadLocal.clear();
                        }
                    });
                });
            } catch (Exception e) {
                //  只是把task add进list而已，还没有到提交那一步。无需执行补偿
                try {
                    tx.rollback(null);
                } catch (Exception ex) {
                    throw new RuntimeException(ex);
                }
                throw new RuntimeException(e);
            }
        });

        return res;
    }
    /**
     *  分表并行
     */
    protected Map<ShardInfo, ShardExecutor.Result> parallelShardTable(Shard anno, IShardTransaction tx, Map<ShardInfo, Object[]> shardInfoMap, Object proxy, Method method) {
        Map<ShardInfo, ShardExecutor.Result> res = new HashMap<>();

        shardInfoMap.forEach((shardInfo, shardArgs) -> {
            String name = String.format("shard_tx_%s_%s_%s_%s", method.getDeclaringClass().getSimpleName(), method.getName(), shardInfo.getShardDB(), shardInfo.getShardTable());
            try {
                //  这里吧result先做好，后面直接用即可
                ShardExecutor.Result r = ShardExecutor.Result.build(shardInfo, false, shardArgs, null, null, -1);
                tx.withTx2PCAsync(name, shardInfo.getShardDB(), r, () -> {
//                    ShardInfoThreadLocal.set(shardInfo);
                    try {
                    	Result _r = ShardExecutor.doInvoke(anno, shardInfo, method, shardArgs);
                    	if (!_r.isSuccess()) {
                    		throw new ShardDBRuntimeException(_r.getE());
                    	}
                        res.put(shardInfo, _r);
                    } finally {
//                        ShardInfoThreadLocal.clear();
                    }
                });
            } catch (Exception e) {
                //  只是把task add进list而已，还没有到提交那一步。无需执行补偿
                try {
                    tx.rollback(null);
                } catch (Exception ex) {
                    throw new RuntimeException(ex);
                }
                throw new RuntimeException(e);
            }
        });
        return res;
    }


    /**
     * 空值映射
     */
    public static ShardTransaction NULL_SHARD_TRANSACTION = new ShardTransaction() {
        public ConcurrentMannerEnums concurrent() {
            return null;
        }
        public Class<? extends IShardCompensation>[] compensation() {
            return null;
        }
        public long timeout() {
            return 0;
        }
        public Class<? extends Annotation> annotationType() {
            return null;
        }
        public boolean strongConsistency() { return false; }
    };
    /**
     *  method与ShardTransacton注解的映射
     */
    protected CachingKVTools<Method, ShardTransaction> shardTransactonMapper = new CachingKVTools<>(NULL_SHARD_TRANSACTION, m -> {
        return m.getAnnotation(ShardTransaction.class);
    });
    /**
     *  事务补偿空值映射
     */
    public static IShardCompensation NULL_COMPENSATION = new IShardCompensation() {
        @Override
        public void doCompensation(List<ShardExecutor.Result> successes, List<ShardExecutor.Result> failies) {
        }
    };
    /**
     *  method与事务补偿映射
     */
    protected CachingKVTools<Method, IShardCompensation> compensationMapper = new CachingKVTools<>(NULL_COMPENSATION, m -> {
        ShardTransaction anno = ShardTxExecutor.this.shardTransactonMapper.get(m);
        if (anno.compensation() == null || anno.compensation().length == 0) {return null;}
        try {
            return TikToolsMyBatisSpringContext.getBeanOrNew(anno.compensation()[0]);
        } catch (InstantiationException | IllegalAccessException e) {
            throw new RuntimeException(e);
        }
    });


    /**
     *  ShardTransactionManager
     */
    protected ShardDataSourceTransactionManager getShardDataSourceTransactionManager() {
        return TikToolsMyBatisSpringContext.getBeanWithAssert(ShardDataSourceTransactionManager.class);
    }
}
